import json
import os
import numpy as np
from tqdm import tqdm

def load_data(file_path):
    """加载JSON行格式的数据"""
    data = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            if line.strip():
                data.append(json.loads(line.strip()))
    return data

def convert_to_sequence_labeling(data):
    """将数据转换为序列标注格式 (BIO标注方案)"""
    processed_data = []
    
    for item in tqdm(data):
        sentence = item['s']
        target = item['ot']
        
        # 查找目标词在句子中的位置
        if target in sentence:
            start = sentence.find(target)
            end = start + len(target)
            
            # 创建BIO标签
            labels = ['O'] * len(sentence)
            
            # 标记B(Begin)
            labels[start] = 'B-TARGET'
            
            # 标记I(Inside)，如果目标长度>1
            for i in range(start + 1, end):
                labels[i] = 'I-TARGET'
                
            processed_data.append({
                'sentence': sentence,
                'chars': list(sentence),
                'labels': labels,
                'target': target
            })
    
    return processed_data

def main():
    # 处理训练、验证和测试数据
    for split in ['train', 'val', 'test']:
        print(f"处理{split}数据...")
        data_path = f'data/{split}.txt'
        
        if os.path.exists(data_path):
            data = load_data(data_path)
            processed_data = convert_to_sequence_labeling(data)
            
            # 保存为处理后的格式
            with open(f'data/{split}_processed.json', 'w', encoding='utf-8') as f:
                json.dump(processed_data, f, ensure_ascii=False, indent=2)
            
            print(f"{split}数据处理完成，共{len(processed_data)}条")

if __name__ == "__main__":
    main()